import os

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain.output_parsers import BooleanOutputParser
from langchain_core.output_parsers import StrOutputParser
from langchain.prompts import ChatPromptTemplate

if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")

    prompt_template = ChatPromptTemplate([
        ("system", "你是一个乐于助人的助手，请根据用户的问题给出回答"),
        ("user", "这是用户的问题:{question},请用 yes 或 no 回答")
    ])
    # bool_qa_chain = prompt_template | model | StrOutputParser()
    bool_qa_chain = prompt_template | model | BooleanOutputParser()
    question = "请问1+1是否大于2？"
    result = bool_qa_chain.invoke(question)
    print(result)
